package com.openworm;
import static java.lang.System.nanoTime;
import static java.lang.System.out;
import static org.bridj.Pointer.allocateFloats;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.List;
import java.util.Random;

import org.bridj.Pointer;
import org.jfree.chart.ChartFactory;
import org.jfree.chart.ChartUtilities;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.data.xy.XYSeries;
import org.jfree.data.xy.XYSeriesCollection;

import com.nativelibs4java.opencl.CLBuffer;
import com.nativelibs4java.opencl.CLContext;
import com.nativelibs4java.opencl.CLEvent;
import com.nativelibs4java.opencl.CLKernel;
import com.nativelibs4java.opencl.CLPlatform.DeviceFeature;
import com.nativelibs4java.opencl.CLProgram;
import com.nativelibs4java.opencl.CLQueue;
import com.nativelibs4java.opencl.JavaCL;
import com.nativelibs4java.opencl.CLMem.Usage;
import com.nativelibs4java.util.IOUtils;

public class AlphaHHKernelTest {

	public static Random randomGenerator = new Random();

	public static float START_TIME = -30;
	public static float END_TIME = 100;

	public static int ELEM_COUNT = 302;

	public static boolean PLOTTING = false;
	public static int SAMPLES = 3;
	
	// conductances
	static float maxG_K = 36;
	static float maxG_Na = 120;
	static float maxG_Leak = (float) 0.3;
	// reverse potentials
	static float E_K = -12;
	static float E_Na = 115;
	static float E_Leak = (float) 10.613;
	// I_ext
	static float I_ext = 10;
	// integration step
	static float dt = (float) 0.01;
	// total time steps
	static int steps = (int) ((int)(END_TIME - START_TIME)/dt);
	
	public static void main(String[] args) throws IOException {
		
		out.println("HH neurons: " + ELEM_COUNT + " - dt: " + dt + " - steps: " + steps);

		CLContext context = JavaCL.createBestContext(DeviceFeature.CPU);
		out.println(context.getDevices()[0].toString());
        CLQueue queue = context.createDefaultQueue();
        
		// Read the program sources and compile them :
		String src = IOUtils.readText(AlphaHHKernelTest.class.getResource("/resources/AlphaHHKernel_Tuning.cl"));
		CLProgram program = context.createProgram(src);
		
		// input buffers declarations
        Pointer<Float> V_in_Ptr = allocateFloats(ELEM_COUNT);
        Pointer<Float> x_n_in_Ptr = allocateFloats(ELEM_COUNT);
        Pointer<Float> x_m_in_Ptr = allocateFloats(ELEM_COUNT);
        Pointer<Float> x_h_in_Ptr = allocateFloats(ELEM_COUNT);

	    initBuffers(V_in_Ptr, x_n_in_Ptr, x_m_in_Ptr, x_h_in_Ptr);
	    
        // Create OpenCL input buffers (using the native memory pointers) :
        CLBuffer<Float> V_in_Buffer = context.createFloatBuffer(Usage.Input, V_in_Ptr);
        CLBuffer<Float> x_n_in_Buffer = context.createFloatBuffer(Usage.Input, x_n_in_Ptr);
        CLBuffer<Float> x_m_in_Buffer = context.createFloatBuffer(Usage.Input, x_m_in_Ptr);
        CLBuffer<Float> x_h_in_Buffer = context.createFloatBuffer(Usage.Input, x_h_in_Ptr);
        CLBuffer<Float> plot_results_Buffer = context.createFloatBuffer(Usage.Output, ELEM_COUNT*steps);

        // save time to track performance
        long compuTime = nanoTime();
        
		// Get and call the kernel
		CLKernel integrateHHStepKernel = program.createKernel("IntegrateHHStep");
		integrateHHStepKernel.setArgs(maxG_K, maxG_Na, maxG_Leak, 
									  E_K, E_Na, E_Leak, I_ext, 
									  dt, steps,
									  V_in_Buffer, x_n_in_Buffer, x_m_in_Buffer, x_h_in_Buffer, plot_results_Buffer,
									  ELEM_COUNT);
		
		int[] globalSizes = new int[] { ELEM_COUNT };
		CLEvent integrateEvt = integrateHHStepKernel.enqueueNDRange(queue, globalSizes);
		// blocks until add_floats finished
		Pointer<Float> outPtr = plot_results_Buffer.read(queue, integrateEvt); 
		
		// print out how long it took
		compuTime = nanoTime() - compuTime;
		out.println("computation took: " + (compuTime/1000000) +"ms");
		
		out.println("results buffer length: " + outPtr.getFloats().length);
		
		/* PLOTTING SETUP */
		if(PLOTTING)
        {
			float[] allFloats = outPtr.getFloats();
			
			// some dictionary for plotting
			Hashtable<Integer, Hashtable<Float, Float>> V_by_t = new Hashtable<Integer, Hashtable<Float, Float>>();
			List<Integer> sampleIndexes = new ArrayList<Integer>();
    		// Generate some random indexes in the 0 .. ELEM_COUNT range
    		for(int i = 0; i < SAMPLES; i++ )
    		{
    			sampleIndexes.add(randomGenerator.nextInt(ELEM_COUNT));
    		}

        	// record some values for plotting
        	Iterator<Integer> itr = sampleIndexes.iterator();
        	while(itr.hasNext())
        	{   
        		Integer index = itr.next();

        		if(!V_by_t.containsKey(index))
        		{
        			V_by_t.put(index, new Hashtable<Float, Float>());
        		}
        		
        		for(int j = 0; j < steps; j++)
        		{
        			V_by_t.get(index).put(new Float(j*dt + START_TIME), allFloats[index + j*ELEM_COUNT]);
        		}
        	}
        	
        	// print some sampled charts to make sure we got fine-looking results.
        	// Plot results
        	Iterator<Integer> iter = sampleIndexes.iterator();
        	while(iter.hasNext())
        	{   
        		XYSeries series = new XYSeries("HH_Graph");

        		Integer index = iter.next();

        		for (int t = 0; t < steps; t++) {
        			series.add(t*dt + START_TIME, V_by_t.get(index).get(t*dt + START_TIME));
        		}

        		// Add the series to your data set
        		XYSeriesCollection dataset = new XYSeriesCollection();
        		dataset.addSeries(series);

        		plotOutput(dataset, index);
        	}
        }
		
		System.out.println("end of HH simulation");
	}

	private static void initBuffers(Pointer<Float> V_in_Ptr, Pointer<Float> x_n_in_Ptr, Pointer<Float> x_m_in_Ptr, Pointer<Float> x_h_in_Ptr) {
		for (int i = 0; i < ELEM_COUNT; i++) {
	    	V_in_Ptr.set(i, (float)-10);
	    	x_n_in_Ptr.set(i, (float)0);
	    	x_m_in_Ptr.set(i, (float)0);
	    	x_h_in_Ptr.set(i, (float)1);
	    }
	}
	
    private static void plotOutput(XYSeriesCollection dataset, int index)
    {
    	// Generate the graph
		JFreeChart chart = ChartFactory.createXYLineChart("HH Chart", "time", "Voltage", dataset, PlotOrientation.VERTICAL, true, true, false);
		try {
			ChartUtilities.saveChartAsJPEG(new File("output/HH_Chart_" + index + ".jpg"), chart, 500, 300);
		} catch (IOException e) {
			System.err.println("Problem occurred creating chart.");
		}
    }
}
